import asyncio
import sys
import argparse
import signal
import re
from fastmcp import FastMCP
from fastmcp import Context
from fastmcp import settings
from typing import Dict, Any, List
from client_manager import ClientManager
from connection_manager import ConnectionManager
from connection import Connection
from config import Config
from logger import Logger
from device import *


# 创建新的MCP实例
mcp = FastMCP(name="Maipu AI Terminal 4T")
config = Config()
client_manager = ClientManager()
logger = Logger("server")
client_lite = 1

async def get_conn_manager(context: Context) -> ConnectionManager:
    """获取连接管理器"""
    # 获取客户端ID

    session_id = "client_lite" if client_lite else context.session_id
    if session_id is None:
        return None
    conn_manager = await client_manager.get_client_connection_manager(session_id)
    return conn_manager

async def close_connection(conn_manager: ConnectionManager, name: str) -> Dict[str, Any]:
    """关闭指定连接
    参数:
    - name: 要关闭的连接名称
    """

    try:
        ret = await conn_manager.close_connection(name)
        # 检查connection_manager返回的是否是错误信息
        if isinstance(ret, dict) and 'error' in ret:
            conn_manager.logger.server_log(f"关闭连接 {name} 失败: {ret['error']}")
            return {
                'success': False,
                'message': ret['error']
            }
        return {
            'success': True,
            'message': f'连接 {name} 已关闭'
        }
    except Exception as e:
        conn_manager.logger.server_log(f'关闭连接 {name} 失败: {str(e)}')
        return {
            'success': False,
            'message': f'关闭连接失败: {str(e)}'
        }

@mcp.tool()
async def send_command(context: Context, device_name: str, command: str) -> Dict[str, Any]:
    """向指定设备发送需要执行的命令;
    MCP会自动连接设备,可以直接发送命令;
    参数:
    - device_name: 设备名称
    - command: 要执行的命令
    """

    conn_manager = await get_conn_manager(context)
    connection = await conn_manager.get_connection(device_name)
    if isinstance(connection, dict):
        return connection
    mode = get_command_mode(command)
    # conn_manager.logger.server_log(f'命令 {command} mode {mode.name}')
    ret, ret_str = await connection.switch_mode(mode)
    if not ret:
        return {
            'success': False,
            'message': ret_str
        }
    conn_manager.logger.server_log(f'命令 {command} mode {mode.name}')
    ret, result = await connection.send_command(command)
    return {
        'success': True if ret else False,
        'message': result
    }

def setup_argparse():
    parser = argparse.ArgumentParser(description='Maipu AI Terminal 4T')

    parser.add_argument('--log_dir', default='./logs', help='日志文件夹路径')
    parser.add_argument('--transport', default='sse', help='传输方式,默认sse')
    parser.add_argument('--host', default='0.0.0.0', help='服务器主机地址,默认0.0.0.0')
    parser.add_argument('--port', type=int, default=19999, help='服务器端口,默认19999')

    return parser

def init(args):
    # 初始化主线程配置和日志
    config.set_log_dir(args.log_dir)
    logger.set_log_dir(args.log_dir)
    logger.register_connection('server', 'localhost', 0)
    logger.server_log('MCP服务器线程已启动')
    logger.server_log(f'日志目录: {config.log_dir}')
    logger.server_log(f'传输方式: {args.transport}')

if __name__ == '__main__':
    parser = setup_argparse()
    args = parser.parse_args()

    init(args)
    # 运行服务器
    if args.transport == "stdio":
        mcp.run(transport=args.transport)
    else:
        mcp.run(transport=args.transport, host=args.host, port=args.port)

    # 关闭所有客户端连接
    client_manager.close_all_clients()
    logger.server_log('服务退出')
